import sys
import getopt
import numpy as np
import mindspore.context as context
from mindspore import Tensor
from mindspore import nn 
import mindspore.ops.operations as P
from mindspore.common import dtype as mstype

opts,_ = getopt.getopt(sys.argv[1:], "d:t:")
opts = dict(opts)
dev_id = 0 if '-d' not in opts else int(opts['-d'])
test_type = 0 if '-t' not in opts else int(opts['-t'])

context.set_context(mode=context.GRAPH_MODE, device_target="GPU", device_id=dev_id)
if test_type > 0:
    context.set_context(enable_graph_kernel=True)
if test_type > 1:
    context.set_context(graph_kernel_flags="--enable_parallel_fusion")

class Net(nn.Cell):
    def __init__(self):
        super(Net, self).__init__()
        self.cast = P.Cast()
        self.sum = P.ReduceSum(keep_dims=True)
        self.addn = P.AddN()

    def construct(self, x1, x2, x3, x4, x5, x6, x7):
        rs1 = self.sum(self.cast(x1, mstype.float32), (0,))
        rs2 = self.sum(self.cast(x2, mstype.float32), (0,))
        rs3 = self.sum(self.cast(x3, mstype.float32), (0,))
        rs4 = self.sum(self.cast(x4, mstype.float32), (0,))
        rs5 = self.sum(self.cast(x5, mstype.float32), (0,))
        rs6 = self.sum(self.cast(x6, mstype.float32), (0,))
        rs7 = self.sum(self.cast(x7, mstype.float32), (0,))
        return self.addn([rs1, rs2, rs3, rs4, rs5, rs6]), rs7 + 0.5

test_shape = [[16, 3072], [16, 21128]]
i0 = Tensor(np.random.normal(0, 1, test_shape[0]).astype(np.float16))
i1 = Tensor(np.random.normal(0, 1, test_shape[0]).astype(np.float16))
i2 = Tensor(np.random.normal(0, 1, test_shape[0]).astype(np.float16))
i3 = Tensor(np.random.normal(0, 1, test_shape[0]).astype(np.float16))
i4 = Tensor(np.random.normal(0, 1, test_shape[0]).astype(np.float16))
i5 = Tensor(np.random.normal(0, 1, test_shape[0]).astype(np.float16))
i6 = Tensor(np.random.normal(0, 1, test_shape[1]).astype(np.float16))

expect = Net()(i0, i1, i2, i3, i4, i5, i6)
